# -*- coding: UTF-8 -*-
import os
import netCDF4 as nc
import numpy as np
from osgeo import gdal, osr, ogr
import glob
from tqdm import tqdm

# os.environ['PROJ_LIB'] = r'./proj'

# 用于逐月数据集

def get_geotransform(nc_file: str, key:str):
    data = nc.Dataset(nc_file)  # 利用.Dataset()读取nc数据
    Lat = data.variables['lat'][:]
    Lon = data.variables['lon'][:]
    data_arr = np.asarray(data.variables[key])  # 属性变量名
    # 影像的左上角&右下角坐标
    Lonmin, Latmax, Lonmax, Latmin = [Lon.min(), Lat.max(), Lon.max(), Lat.min()]
    # Lonmin, Latmax, Lonmax, Latmin
    # 分辨率计算
    Num_lat = len(Lat)
    Num_lon = len(Lon)
    Lat_res = (Latmax - Latmin) / (float(Num_lat) - 1)
    Lon_res = (Lonmax - Lonmin) / (float(Num_lon) - 1)
    # 设置影像的显示范围
    # Lat_re前需要添加负号
    nodata =data.variables[key].missing_value  # netcdf文件的缺失值
    geotransform = (Lonmin, Lon_res, 0.0, Latmax, 0.0, -Lat_res)
    mask = np.where(data_arr[0,:,:] == nodata, 0, 1)  # nodata区域为 0, 多波段只读取一个波段
    return geotransform, nodata, mask

def arr2tif(data_arr, out_tif_name: str, geotransform,nodata):
    # 创建tif文件
    driver = gdal.GetDriverByName('GTiff')
    out_tif = driver.Create(out_tif_name, data_arr.shape[1],data_arr.shape[0],1, gdal.GDT_Float32)
    # driver.Create(filename, width, height, bands, DataType.GDT_Byte, null)
    out_tif.SetGeoTransform(geotransform)
    # 定义投影
    prj = osr.SpatialReference()
    prj.ImportFromEPSG(4326)
    out_tif.SetProjection(prj.ExportToWkt())
    out_tif.GetRasterBand(1).SetNoDataValue(nodata)
    # 数据导出
    out_tif.GetRasterBand(1).WriteArray(data_arr)  # 将数据写入内存
    out_tif.FlushCache()  # 将数据写入到硬盘
    out_tif = None  # 关闭tif文件

    return out_tif_name


if __name__ == "__main__":

    # #测试数据信息，每一个nc文件是逐月降水、均温数据
    # #中国1km分辨率逐月降水量数据集（1901 - 2023）
    # #https://data.tpdc.ac.cn/zh-hans/data/faae7605-a0f2-4d18-b28f-5cee413766a2
    #中国1km分辨率逐月平均气温数据集（1901-2023）
    #https://data.tpdc.ac.cn/zh-hans/data/71ab4677-b66c-4fd1-a004-b2a541c4d5bf

    Input_pre = r"D:\Users\Jiudian\Desktop\pre"
    Input_tmp = r"D:\Users\Jiudian\Desktop\tmp"
    Output_folder = r"D:\Users\Jiudian\Desktop\Rh"
    start_year = 2000
    end_year = 2023
    key1 = 'pre'  # 均温是tmp
    key2 = 'tmp'  # 降水是pre
    band_dimension = 12  # 时间维度，12个月

    if not os.path.exists(Output_folder):  # 如果目录不存在就返回False，创建该文件夹
        os.mkdir(Output_folder)

    # 读取所有数据
    pre_list = sorted(glob.glob(os.path.join(Input_pre, f"*.nc")))
    tmp_list = sorted(glob.glob(os.path.join(Input_tmp, f"*.nc")))

    geotran, nodata, mask= get_geotransform(pre_list[1], key1)
    _mask = np.where(mask == 0, 1, 0)  # mask的反
    if len(pre_list) == len(tmp_list): # 如果两个文件夹的nc文件数量相等，进行计算
        # 建一个循环，每个循环分别读取月降水和月均温，计算逐月Rh，最后sum计算年Rh
        for k in tqdm(range(start_year, end_year+1)):  # 循环每一个nc文件，内有12个波段/dims
            j = k-start_year
            pre_arr = np.asarray(nc.Dataset(pre_list[j]).variables[key1]) * 0.01  # 比例转换 0.1mm 计算公式单位是cm
            tmp_arr = np.asarray(nc.Dataset(tmp_list[j]).variables[key2]) * 0.1   # 比例转换 0.1deg
            M = np.zeros((pre_arr.shape[1], pre_arr.shape[2]))  # 创建一个数组 m*n
            Rh = np.zeros((band_dimension, pre_arr.shape[1], pre_arr.shape[2]))  # 创建一个数组 time*m*n
            for i in range(band_dimension):  # 计算逐月Rs与Rh
                Rs = 1.25 * np.exp(0.05452 * tmp_arr[i, :, :]) * (pre_arr[i, :, :] / (4.259 + pre_arr[i, :, :]) )
                Rh[i, :, :] = np.exp(0.22 + 0.87 * np.log(Rs))
            M = Rh.sum(axis=0) * mask + _mask * nodata  # 累积年Rh，把nodata数据添加上
            arr2tif(M, os.path.join(Output_folder, "Rh_"f"{k}.tif"), geotran, nodata)  # 保存年Rh为tif
            print(f"{k}年Rh计算完成！")
    else:
        print("降水、均温数据集数量不一致，请检查！")
